# -*- coding: utf-8 -*-
import codecs
import sys,os
#import wikipedia
import cPickle
import time
#import Levenshtein

 
def extract_xml_attrnode(line):
    tmp1=line.split('>')
    idx = tmp1[0].find('=')
    tmp_id = tmp1[0][idx+1:].replace('"','').strip()
    idx = tmp1[1].find('<')
    tmp_name = tmp1[1][0:idx]
    return (tmp_id,tmp_name)

def extract_xml_node(line):
    tmp1=line.split('>')
    idx = tmp1[1].find('<')
    tmp_name = tmp1[1][0:idx]
    return tmp_name

def extract_xml_att(line):
    tmp1=line.split('>')
    idx = tmp1[0].find('=')
    tmp_id = tmp1[0][idx+1:].replace('"','').strip()
    return tmp_id

def readSample(kb_file):
    fid = codecs.open(kb_file, 'r', encoding='utf-8')
    raw=fid.read()
    fid.close()
    lines = raw.split('\n')
    
    gquery_set = []
    gquery_id = ''
    gquery_content = ''
    gquery_entity_set= []
    gtmp_id = ''
    gtmp_name = ''
    gtmp_position = ['','']
    gtmp_kbid= ''
    for ilx in range(0,len(lines)):
        line = lines[ilx].strip()
        if line[0:6]=='<?xml ':
            continue
        
        if line[0:6]=='<weibo':
            gquery_id = extract_xml_att(line)
            
        if line[0:8]=='</weibo>':
            query = (gquery_id,gquery_content,gquery_entity_set)
            gquery_set.append(query)
            gquery_id = ''
            gquery_content = ''
            gquery_entity_set= []
        
        if line[0:5]=='<name':
            (gtmp_id,gtmp_name)=extract_xml_attrnode(line)
            
        if line[0:12]=='<startoffset':
            (tmp_id1,tmp_off1)=extract_xml_attrnode(line)
            assert(tmp_id1==gtmp_id)
            gtmp_position[0] = tmp_off1
        
        if line[0:10]=='<endoffset':
            (tmp_id1,tmp_off1)=extract_xml_attrnode(line)
            assert(tmp_id1==gtmp_id)
            gtmp_position[1] = tmp_off1
        
        if line[0:3]=='<kb':
            (tmp_id1,gtmp_kbid)=extract_xml_attrnode(line)
            assert(tmp_id1==gtmp_id)
            entity = (gtmp_id,gtmp_name,gtmp_position,gtmp_kbid)
            gquery_entity_set.append(entity)
            
            gtmp_id = ''
            gtmp_name = ''
            gtmp_position = ['','']
            gtmp_kbid= ''
            
        if line[0:9]=='<content>':
            gquery_content = extract_xml_node(line)
            
    return gquery_set
 

#wikipedia.set_lang("zh")
BASE_DATA_PATH  = 'D:\\data\\Shared_Task_NLPCC14\\'
BASE_KB_PATH = 'E:\\desktop\\wu-request\\NLPCC 2014 Shared Tasks Guidelines\\Chinese Entity Linking  SAMPLE DATA NLPCC2014_EL_sample\\'

BASELINE_STEP = 3

'step0. 加载sample_query_2014和PKBase_key_title'
if BASELINE_STEP<=5:
    kb_path = BASE_KB_PATH
    sp_file = kb_path+'weiboAutoTag_6.txt'
    query_set = readSample(sp_file)
    kb_file = kb_path + 'PKBase_key_title.txt'
    #(map_key2name,map_name2key,name_list) = load_KBaseKey(kb_file)
    print 'Load sample and key finished!'
    

find_qset=[]
for qs in query_set:
    cc=qs[2]
    for c in cc:
        if len(c[1])==2:
            print c[1].encode('gbk','ignore'),qs[0].encode('gbk','ignore')
            find_qset.append(qs)
            break
    
    
    
    